#include <boost/asio.hpp>
#include <boost/beast.hpp>
#include <boost/beast/websocket.hpp>
#include <boost/asio/ssl.hpp> // For SSL/TLS
#include <boost/beast/ssl.hpp> // <--- Add this line
#include <iostream>
#include <memory>
#include <string>
#include <deque> // For message queue
#include <Poco/SHA1Engine.h>
#include <Poco/Crypto/DigestEngine.h>
#include <Poco/Base64Encoder.h>
#include "AppDirectory.h"
#include "KedisConfig.h"
// Define aliases for convenience
namespace asio = boost::asio;
namespace beast = boost::beast;
namespace http = beast::http;
namespace websocket = beast::websocket;
namespace ssl = asio::ssl;
using tcp = asio::ip::tcp;

// // --- Mocking External Dependencies ---
// // In a real application, these would be passed or resolved properly.
// // const std::string BACKEND_HOST_WS = "v.host";
// // const unsigned short BACKEND_PORT_WS = 8081;
// // const std::string BACKEND_HOST_HTTP = "v.host";
// // const unsigned short BACKEND_PORT_HTTP = 8080;

// // A simple mock for an HTTP request
// struct MockHttpRequest {
//     http::verb method_ = http::verb::get;
//     std::string target_ = "/echo";
//     unsigned int version_ = 11; // HTTP/1.1
//     std::string websocket_key_ = "dGhlIHNhbXBsZSBub25jZQ=="; // A dummy key
//     std::string upgrade_header_ = "websocket";
//     std::string connection_header_ = "Upgrade";

//     http::verb method() const { return method_; }
//     std::string target() const { return target_; }
//     unsigned int version() const { return version_; }
//     std::string operator[](http::field field) const {
//         if (field == http::field::sec_websocket_key) return websocket_key_;
//         if (field == http::field::upgrade) return upgrade_header_;
//         if (field == http::field::connection) return connection_header_;
//         return "";
//     }
// };
// // --- End Mocking ---

// Forward declarations
class session; // Declare session class

// Type alias for the client's SSL/TLS stream
using ssl_stream = ssl::stream<tcp::socket>;

// The main session class, inheriting from enable_shared_from_this for lifetime management
class session : public std::enable_shared_from_this<session> {
public:
    // Constructor
    session(asio::io_context& ioc, ssl::context& ssl_ctx, tcp::socket socket, KedisConfig& config)
        : ioc_(ioc),
          ssl_ctx_(ssl_ctx),
          stream_(std::move(socket), ssl_ctx_), // Initialize SSL stream with the accepted socket
          resolver_(ioc),
          backend_socket_(ioc), // Initialize backend TCP socket
          config_(config)
    {
        std::cout << getFormattedTime() << "session::session() - New session created.\n";
    }

    // Entry point for starting the session
    void start() {
        std::cout << getFormattedTime() << "session::start() - Starting SSL handshake...\n";
        auto self = shared_from_this(); // Keep session alive

        // Perform SSL handshake with the client
        stream_.async_handshake(ssl::stream_base::server,
            [self](beast::error_code ec) {
                if (ec) {
                    std::cerr << self->getFormattedTime() << "SSL handshake failed: " << ec.message() << "\n";
                    self->close_connections(); // Close if handshake fails
                    return;
                }
                std::cout << self->getFormattedTime() << "SSL handshake success.\n";
                self->read_request(); // Proceed to read HTTP request
            }
        );
    }

    // std::string base64_sha1(const std::string& input) {
    //     // 1. Calculate SHA-1 hash
    //     Poco::SHA1Engine sha1;
    //     sha1.update(input);
    //     std::vector<unsigned char> digest = sha1.digest(); // Get digest directly
    
    //     // Convert hash to string
    //     std::string sha1_hash_str(digest.begin(), digest.end());
    
    //     // 2. Base64 encode
    //     std::stringstream ss;
    //     Poco::Base64Encoder encoder(ss);
    //     encoder.write(sha1_hash_str.c_str(), sha1_hash_str.size());
    //     encoder.close(); // Important to flush and finalize the encoding
    //     return ss.str();
    // }

private:
    asio::io_context& ioc_;
    ssl::context& ssl_ctx_;
    ssl_stream stream_; // Client's SSL stream
    tcp::socket backend_socket_; // Backend plain TCP socket (used for initial connection)
    tcp::resolver resolver_; // Resolver for backend addresses

    // WebSocket streams (shared_ptr for optionality and lifetime)
    std::shared_ptr<websocket::stream<ssl_stream>> client_ws_;
    std::shared_ptr<websocket::stream<tcp::socket>> backend_ws_;

    http::request<http::string_body> req_; // Client HTTP request
    // MockHttpRequest mock_req_; // Using mock for simplicity in this example

    // Queue for messages from client to backend
    std::deque<std::string> client_to_backend_queue_;
    // Queue for messages from backend to client
    std::deque<std::string> backend_to_client_queue_;
    beast::flat_buffer buffer_; // <--- Add this line for the read buffer
    KedisConfig& config_;
    // --- Helper for formatted time logging ---
    std::string getFormattedTime() {
        auto now = std::chrono::system_clock::now();
        auto in_time_t = std::chrono::system_clock::to_time_t(now);
        std::stringstream ss;
        ss << std::put_time(std::localtime(&in_time_t), "%Y-%m-%d %H:%M:%S.");
        auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch()) % 1000;
        ss << std::setfill('0') << std::setw(3) << ms.count() << " ";
        return ss.str();
    }
    // 函数用于获取当前时间并格式化为 "YYYY-MM-DD HH:MM:SS.ms" 字符串
    // std::string getFormattedTime() {
    //     // 获取当前系统时间点
    //     auto now = std::chrono::system_clock::now();

    //     // 将时间点转换为 time_t，以便使用 localtime
    //     std::time_t now_c = std::chrono::system_clock::to_time_t(now);

    //     // 获取毫秒部分
    //     auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch()) % 1000;

    //     // 将 time_t 转换为 tm 结构，用于格式化
    //     // 注意：localtime 不是线程安全的，在多线程环境下应使用 localtime_s (Windows) 或 localtime_r (POSIX)
    //     // Use thread-safe version if available, e.g., std::put_time is safe with its own tm struct
    //     // For C++20, std::chrono::format is even better.
    //     std::tm tm_buf;
    //     #ifdef _WIN32
    //         localtime_s(&tm_buf, &now_c); // Windows
    //     #else
    //         localtime_r(&now_c, &tm_buf); // POSIX
    //     #endif

    //     // 使用 put_time 格式化日期和时间部分
    //     std::ostringstream oss;
    //     oss << std::put_time(&tm_buf, "%Y-%m-%d %H:%M:%S");

    //     // 添加毫秒部分
    //     oss << '.' << std::setfill('0') << std::setw(3) << ms.count();

    //     return oss.str();
    // }

    // --- Core Logic Functions (Callback Style) ---

    // 1. Read HTTP Request from Client
    void read_request() {
        std::cout << getFormattedTime() << "session::read_request() - Reading HTTP request...\n";
        auto self = shared_from_this();

        // Clear the buffer before each read to prevent old data issues
        // and ensure it's ready for new incoming data.
        buffer_.clear();

        // Use beast::http::async_read to read the HTTP request
        http::async_read(stream_, buffer_, req_,
            [self](beast::error_code ec, std::size_t bytes_transferred) {
                if (ec) {
                    std::cerr << self->getFormattedTime() << "Error reading HTTP request: " << ec.message() << "\n";
                    self->close_connections();
                    return;
                }
                std::cout << self->getFormattedTime() << "Read request: method=" << self->req_.method()
                          << ", target=" << self->req_.target() << "\n";

                // Check if it's a WebSocket upgrade request
                if (websocket::is_upgrade(self->req_)) {
                    std::cout << self->getFormattedTime() << "Request is a WebSocket upgrade.\n";
                    self->handle_websocket_upgrade();
                } else {
                    self->forward_http_request();
                    // std::cerr << self->getFormattedTime() << "Non-WebSocket request received. Not supported by this proxy.\n";
                    // self->close_connections(); // Close for unsupported requests
                }
            }
        );
    }

    void forward_http_request() {
        // Connect to backend HTTP server
        // This part is also synchronous. Consider converting to async if this proxy needs to handle many concurrent HTTP requests.
        try {
            resolve_backend_http();
        } catch (const std::exception& e) {
            std::cerr << getFormattedTime() << " Exception in forward_http_request: " << e.what() << "\n";
            send_bad_gateway_http_response();
        }
    }

    void read_http_response() {
        auto self = shared_from_this(); // Ensure session lifetime
        std::cout << getFormattedTime() << "read_http_response() - Reading backend HTTP response.\n";

        // Use a shared_ptr to manage the buffer's lifetime
        auto backend_buffer = std::make_shared<beast::flat_buffer>();
        auto res = std::make_shared<http::response<beast::http::dynamic_body>>();

        if (!self->backend_socket_.is_open()) {
            std::cerr << self->getFormattedTime() << "Error: Backend socket is not open before async_read_http_response!\n";
            // Handle this error appropriately, perhaps send 502 Bad Gateway
            self->send_bad_gateway_http_response();
            return;
        }

        std::cout << getFormattedTime() << "read_http_response before async_read...\n";
        http::async_read(self->backend_socket_, *backend_buffer, *res,
            [self, backend_buffer, res](beast::error_code ec_read_res, std::size_t bytes_transferred) {
                std::cout << self->getFormattedTime() << "session::read_http_response async_read completed, bytes transferred: " << bytes_transferred << "\n";
                std::cout << self->getFormattedTime() << "Backend response status: " << res->result_int() << "\n";
                // std::cout << self->getFormattedTime() << "Backend response Content-Length: " << res->has_content_length() ? std::to_string(res->content_length().value()) : "N/A" << "\n";
                std::cout << self->getFormattedTime() << "Backend response body size (dynamic_body): " << res->body().size() << "\n";

                std::cout << self->getFormattedTime() << "Raw buffer data after async_read:\n";
                // This will print the raw bytes Beast read into the buffer
                // It's helpful to see if the image data is actually there, but not parsed into the body
                std::string raw_data = boost::beast::buffers_to_string(backend_buffer->data());
                std::cout << raw_data.substr(0, std::min((size_t)200, raw_data.length())) << "...\n"; // Print first 200 chars

                // 尝试打印body的前几个字节，用于调试
                if (res->body().size() > 0) {
                    auto data_part = boost::beast::buffers_to_string(res->body().data());
                    std::cout << self->getFormattedTime() << "Backend body starts with: " << data_part.substr(0, std::min((size_t)50, data_part.length())) << "...\n";
                }
                if (ec_read_res && ec_read_res != http::error::end_of_stream) {
                    std::cerr << self->getFormattedTime() << " Error reading HTTP response from backend server: " << ec_read_res.message() << "\n";
                    // Send an error response to the client
                    auto err_res = std::make_shared<http::response<http::string_body>>(
                        http::status::internal_server_error, self->req_.version());
                    err_res->set(http::field::server, "Boost.Beast Reverse Proxy");
                    err_res->body() = "Failed to read response from backend.";
                    err_res->prepare_payload();

                    // Asynchronously write the error response to the client
                    http::async_write(self->stream_, *err_res,
                        [self](beast::error_code ignored_ec, std::size_t /*bytes_written*/) {
                            if (ignored_ec) {
                                std::cerr << self->getFormattedTime() << " Error writing error response to client: " << ignored_ec.message() << "\n";
                            }
                            self->close_connections_http(); // Close connections after error
                        }
                    );
                    return;
                }

                // Check if the backend response indicates connection closure
                bool backend_keep_alive = res->keep_alive();
                if (!backend_keep_alive || ec_read_res == http::error::end_of_stream) {
                    std::cout << self->getFormattedTime() << " Backend closed connection or sent Connection: close.\n";
                    beast::error_code ec;
                    self->backend_socket_.shutdown(tcp::socket::shutdown_both, ec);
                    self->backend_socket_.close(ec);
                }

                // Send the response to the client
                self->send_http_response_to_client(res);
            }
        );
    }

    std::string get_mime_type(const std::string& path) {
        auto const ext = [&path] {
            auto const pos = path.rfind(".");
            if (pos == std::string::npos)
                return std::string{};
            return path.substr(pos);
        }();
        
        if (ext == ".htm")  return "text/html";
        if (ext == ".html") return "text/html";
        if (ext == ".php")  return "text/html";
        if (ext == ".css")  return "text/css";
        if (ext == ".txt")  return "text/plain";
        if (ext == ".js")   return "application/javascript";
        if (ext == ".json") return "application/json";
        if (ext == ".xml")  return "application/xml";
        if (ext == ".swf")  return "application/x-shockwave-flash";
        if (ext == ".flv")  return "video/x-flv";
        if (ext == ".png")  return "image/png";
        if (ext == ".jpe")  return "image/jpeg";
        if (ext == ".jpeg") return "image/jpeg";
        if (ext == ".jpg")  return "image/jpeg";
        if (ext == ".gif")  return "image/gif";
        if (ext == ".bmp")  return "image/bmp";
        if (ext == ".ico")  return "image/vnd.microsoft.icon";
        if (ext == ".tiff") return "image/tiff";
        if (ext == ".tif")  return "image/tiff";
        if (ext == ".svg")  return "image/svg+xml";
        if (ext == ".svgz") return "image/svg+xml";
        return "application/octet-stream";
    }

    void send_http_response_to_client(std::shared_ptr<http::response<beast::http::dynamic_body>> res) {
        beast::error_code ec_write_client;

        auto self = shared_from_this();
        std::cout << getFormattedTime() << "session::send_http_response() - Sending HTTP response to client...\n";
        // 设置正确的MIME类型（如果后端没有设置）
        if (!res->count(http::field::content_type)) {
            std::string path = self->req_.target();
            res->set(http::field::content_type, get_mime_type(path));
        }

        http::async_write(self->stream_, *res,
            [self, res](beast::error_code ec_write_client, std::size_t bytes_written) {
                if (ec_write_client) {
                    std::cerr << self->getFormattedTime() << " Error writing HTTP response to client: " << ec_write_client.message() << "\n";
                    self->close_connections_http();
                    return;
                }
                std::cout << self->getFormattedTime() << " Successfully sent " << bytes_written << " bytes to client.\n";
                // Check if the client connection should remain open (based on response headers)
                bool client_keep_alive = res->keep_alive();
                if (!client_keep_alive) {
                    std::cout << self->getFormattedTime() << " Client connection closed (Connection: close).\n";
                    self->close_connections_http();
                    return;
                }

                // Continue reading the next request if keep-alive is enabled
                if (self->stream_.lowest_layer().is_open()) {
                    std::cout << self->getFormattedTime() << " Connection is keep-alive, awaiting next request...\n";
                    self->read_request();
                } else {
                    std::cout << self->getFormattedTime() << " Client connection closed unexpectedly.\n";
                    self->close_connections_http();
                }
            }
        );
        
        // 关闭 SSL 连接
        // For WebSocket, rely on close_websocket_connections() instead of shutting down stream_
        // if (!client_ws_) {
        //     // Only shutdown SSL stream for non-WebSocket connections
        //     stream_.shutdown(ignored);
        //     if (ignored && ignored != asio::error::eof) {
        //         std::cerr << "Error during SSL shutdown: " << ignored.message() << "\n";
        //     }
        // }
    }

    void send_bad_gateway_http_response() {
        http::response<http::string_body> res{http::status::bad_gateway, req_.version()};
        res.set(http::field::server, "Boost.Beast Reverse Proxy");
        res.body() = "Failed to connect to HTTP backend.";
        res.prepare_payload();
        beast::error_code ec_write;
        http::write(stream_, res, ec_write); // Blocking write
        // Error handling for this write is omitted for brevity
    }

    // 2. Handle WebSocket Upgrade Request (Client Side)
    void handle_websocket_upgrade() {
        std::cout << getFormattedTime() << "session::handle_websocket_upgrade()...\n";
        auto self = shared_from_this();

        // Initialize client WebSocket stream using the existing SSL stream
        client_ws_ = std::make_shared<websocket::stream<ssl_stream>>(std::move(stream_));

        // Set up control frame callback for client (for logging or custom handling)
        client_ws_->control_callback(
            [self](beast::websocket::frame_type kind, std::string_view payload) {
                std::cout << self->getFormattedTime() << " Client Control Frame - Type:" << (int)kind << ", Payload Size:" << payload.size() << "\n";
                if (kind == websocket::frame_type::close) {
                    std::cout << self->getFormattedTime() << " Client sent Close frame. Propagating closure.\n";
                    self->close_connections(); // Initiate closure
                }
            });

        // Perform WebSocket handshake with the client
        // Note: async_accept for client_ws_
        client_ws_->async_accept(req_, // Use mock_req_ for the handshake
            [self](beast::error_code ec_client_accept) {
                if (ec_client_accept) {
                    std::cerr << self->getFormattedTime() << "Error accepting client WebSocket: " << ec_client_accept.message() << "\n";
                    self->close_connections();
                    return;
                }
                std::cout << self->getFormattedTime() << "Client WebSocket handshake succeeded.\n";
                self->resolve_backend_websocket(); // Proceed to resolve backend
            }
        );
    }

    // 3. Resolve Backend Server Address
    void resolve_backend_websocket() {
        std::string BACKEND_HOST_WS = config_.get("backend_host_ws");
        std::string BACKEND_PORT_WS = config_.get("backend_port_ws");
        std::cout << getFormattedTime() << "Resolving backend: " << BACKEND_HOST_WS << ":" << BACKEND_PORT_WS << "...\n";
        auto self = shared_from_this();

        resolver_.async_resolve(BACKEND_HOST_WS, BACKEND_PORT_WS,
            [self](beast::error_code ec_resolve, tcp::resolver::results_type results) {
                if (ec_resolve) {
                    std::cerr << self->getFormattedTime() << "Resolution error: " << ec_resolve.message() << std::endl;
                    self->close_connections();
                    return;
                }
                std::cout << self->getFormattedTime() << "Resolve backend success. Attempting to connect...\n";
                self->connect_backend_websocket(results); // Proceed to connect to backend
            }
        );
    }

    void resolve_backend_http() {
        std::string BACKEND_HOST_HTTP = config_.get("backend_host_http");
        std::string BACKEND_PORT_HTTP = config_.get("backend_port_http");
        std::cout << getFormattedTime() << "Resolving backend: " << BACKEND_HOST_HTTP << ":" << BACKEND_PORT_HTTP << "...\n";
        auto self = shared_from_this();

        resolver_.async_resolve(BACKEND_HOST_HTTP, BACKEND_PORT_HTTP,
            [self](beast::error_code ec_resolve, tcp::resolver::results_type results) {
                if (ec_resolve) {
                    std::cerr << self->getFormattedTime() << "Resolution error: " << ec_resolve.message() << std::endl;
                    self->close_connections();
                    return;
                }
                std::cout << self->getFormattedTime() << "Resolve backend success. Attempting to connect...\n";
                self->connect_backend_http(results); // Proceed to connect to backend
            }
        );
    }

    void connect_backend_http(tcp::resolver::results_type& results) {
        std::cout << getFormattedTime() << "Connecting to backend...\n";
        auto self = shared_from_this();
        #if defined(__GNUC__)
        asio::async_connect(backend_socket_, results, // <--- Corrected
            [self](beast::error_code ec_connect, tcp::endpoint connected_endpoint) {
                if (ec_connect) {
                    std::cerr << self->getFormattedTime() << "Error connecting to backend Http server: " << ec_connect.message() << "\n";
                    self->close_connections();
                    return;
                }

                std::cout << self->getFormattedTime() << "Successfully connected to http backend: "
                          << connected_endpoint.address() << ":" << connected_endpoint.port() << "\n";

                self->send_request_to_backend_http(); // Proceed to backend WebSocket handshake
            }
        );
        #else
        // Change this line:
        asio::async_connect(backend_socket_, results.begin(), results.end(), // <--- Corrected
            [self](beast::error_code ec_connect, tcp::resolver::iterator connected_it) {
                if (ec_connect) {
                    std::cerr << self->getFormattedTime() << "Error connecting to backend Http server: " << ec_connect.message() << "\n";
                    self->close_connections();
                    return;
                }
                tcp::endpoint connected_endpoint = *connected_it;
                std::cout << self->getFormattedTime() << "Successfully connected to http backend: "
                          << connected_endpoint.address() << ":" << connected_endpoint.port() << "\n";

                self->send_request_to_backend_http(); // Proceed to backend WebSocket handshake
            }
        );
        #endif
    }

    // void close_connections_http()
    // {
    //     // 关闭后端 socket
    //     beast::error_code ignored;
    //     backend_socket_.shutdown(asio::ip::tcp::socket::shutdown_both, ignored);            
    //     backend_socket_.close(); // Close backend socket

    //     // For non-WebSocket HTTP, also gracefully shutdown the client SSL stream
    //     // This is crucial to release the client connection.
    //     stream_.shutdown(ignored); 
    //     if (ignored && ignored != asio::error::eof && ignored != asio::error::operation_aborted) {
    //         std::cerr << getFormattedTime() << " Error during client SSL shutdown: " << ignored.message() << "\n";
    //     }
    // }

    void close_connections_http() {
        beast::error_code ec;

        if (backend_socket_.is_open()) {
            std::cout << getFormattedTime() << "Closing underlying backend TCP socket...\n";
            backend_socket_.shutdown(tcp::socket::shutdown_both, ec);
            backend_socket_.close(ec);
            if (ec && ec != asio::error::eof && ec != asio::error::operation_aborted) {
                std::cerr << getFormattedTime() << " Error closing backend socket: " << ec.message() << "\n";
            }
        }

        if (stream_.lowest_layer().is_open()) {
            std::cout << getFormattedTime() << "Closing underlying client SSL stream...\n";
            stream_.lowest_layer().shutdown(tcp::socket::shutdown_both, ec);
            stream_.lowest_layer().close(ec);
            if (ec && ec != asio::error::eof && ec != asio::error::operation_aborted) {
                std::cerr << getFormattedTime() << " Error closing client SSL stream: " << ec.message() << "\n";
            }
        }
    }

    void send_request_to_backend_http()
    {
        std::string BACKEND_HOST_HTTP = config_.get("backend_host_http");
        std::string BACKEND_PORT_HTTP = config_.get("backend_port_http");
        auto self = shared_from_this(); // Ensure session lifetime
        req_.set(http::field::host, BACKEND_HOST_HTTP + ":" + BACKEND_PORT_HTTP);
        
        std::cout << getFormattedTime() << "Writing HTTP request to backend server...\n";
        http::async_write(backend_socket_, req_,
            [self](beast::error_code ec_write_req, std::size_t bytes_transferred) {
                if (ec_write_req) {
                    std::cerr << self->getFormattedTime() << "Error writing HTTP request to backend server: " << ec_write_req.message() << "\n";
                    self->close_connections(); // Close connections on error
                    return;
                }
                std::cout << self->getFormattedTime() << "Successfully wrote " << bytes_transferred << " bytes to backend server.\n";
                // Proceed to the next step, e.g., reading the response
                self->read_http_response(); // Placeholder for next action
            }
        );
    }

    // 4. Connect to Backend Server
    void connect_backend_websocket(tcp::resolver::results_type& results) {
        std::cout << getFormattedTime() << "Connecting to backend...\n";
        auto self = shared_from_this();

        // Change this line:
        #if defined(__GNUC__)
        asio::async_connect(backend_socket_, results, // <--- Corrected
            [self](beast::error_code ec_connect, tcp::endpoint connected_endpoint) {
                if (ec_connect) {
                    std::cerr << self->getFormattedTime() << "Error connecting to backend WebSocket server: " << ec_connect.message() << "\n";
                    self->close_connections();
                    return;
                }

                std::cout << self->getFormattedTime() << "Successfully connected to backend: "
                          << connected_endpoint.address() << ":" << connected_endpoint.port() << "\n";

                self->perform_backend_websocket_handshake(); // Proceed to backend WebSocket handshake
            }
        );
        #else
        asio::async_connect(backend_socket_, results.begin(), results.end(), // <--- Corrected
            [self](beast::error_code ec_connect, tcp::resolver::iterator connected_it) {
                if (ec_connect) {
                    std::cerr << self->getFormattedTime() << "Error connecting to backend WebSocket server: " << ec_connect.message() << "\n";
                    self->close_connections();
                    return;
                }
                tcp::endpoint connected_endpoint = *connected_it;
                std::cout << self->getFormattedTime() << "Successfully connected to backend: "
                          << connected_endpoint.address() << ":" << connected_endpoint.port() << "\n";

                self->perform_backend_websocket_handshake(); // Proceed to backend WebSocket handshake
            }
        );
        #endif
    }

    // 5. Perform Backend WebSocket Handshake
    void perform_backend_websocket_handshake() {
        std::cout << getFormattedTime() << "Performing backend WebSocket handshake...\n";
        auto self = shared_from_this();
        std::string BACKEND_HOST_WS = config_.get("backend_host_ws");
        std::string BACKEND_PORT_WS = config_.get("backend_port_ws");
        // Initialize backend WebSocket stream using the connected TCP socket
        backend_ws_ = std::make_shared<websocket::stream<tcp::socket>>(std::move(backend_socket_));

        // Set up control frame callback for backend
        backend_ws_->control_callback(
            [self](beast::websocket::frame_type kind, std::string_view payload) {
                std::cout << self->getFormattedTime() << " Backend Control Frame - Type:" << (int)kind << ", Payload Size:" << payload.size() << "\n";
                if (kind == websocket::frame_type::close) {
                    std::cout << self->getFormattedTime() << " Backend sent Close frame. Propagating closure.\n";
                    self->close_connections(); // Initiate closure
                }
            });

        // Perform WebSocket handshake with the backend
        // Note: async_handshake for backend_ws_
        backend_ws_->async_handshake(BACKEND_HOST_WS, req_.target(), // Use mock_req_.target() for backend path
            [self](beast::error_code ec_backend_handshake) {
                if (ec_backend_handshake) {
                    std::cerr << self->getFormattedTime() << " Error performing backend WebSocket handshake: " << ec_backend_handshake.message() << "\n";
                    self->close_connections();
                    return;
                }
                std::cout << self->getFormattedTime() << " Backend WebSocket handshake succeeded.\n";
                self->start_forwarding_loops_websocket(); // Proceed to start data forwarding
            }
        );
    }

    // 6. Start Bidirectional Data Forwarding Loops
    void start_forwarding_loops_websocket() {
        std::cout << getFormattedTime() << "Starting concurrent WebSocket data forwarding loops...\n";
        // In callback style, these are typically self-chaining async operations
        // that call themselves in their completion handlers.
        read_from_client_websocket();
        read_from_backend_websocket();
    }

    // Read from client and write to backend
    void read_from_client_websocket() {
        auto self = shared_from_this();
        if (!client_ws_ || !client_ws_->is_open() || !backend_ws_ || !backend_ws_->is_open()) {
            std::cout << getFormattedTime() << " Client or Backend WebSocket closed in read_from_client. Exiting loop.\n";
            close_connections();
            return;
        }

        // beast::flat_buffer buffer;
        auto buffer = std::make_shared<beast::flat_buffer>();
        std::cout << getFormattedTime() << " Client -> Backend: Awaiting async_read from client...\n";
        client_ws_->async_read(*buffer,
            [self, buffer](beast::error_code ec_read, std::size_t bytes_transferred) mutable { // mutable for buffer
                if (ec_read == websocket::error::closed || ec_read == asio::error::eof || ec_read == asio::error::operation_aborted) {
                    std::cout << self->getFormattedTime() << " Client WebSocket closed or aborted in read_from_client. " << ec_read.message() << "\n";
                    self->close_connections();
                    return;
                }
                if (ec_read) {
                    std::cerr << self->getFormattedTime() << " Error reading from client: " << ec_read.message() << "\n";
                    self->close_connections();
                    return;
                }
                // std::string s = beast::buffers_to_string(buffer.data());
                std::cout << self->getFormattedTime() << " Client -> Backend: Read " << bytes_transferred
                            << " bytes. Is Text: " << self->client_ws_->got_text() << ", Is Binary: " << self->client_ws_->got_binary() << "\n";

                // Add this line to see the actual size of the payload available in the buffer
                std::string str_received = beast::buffers_to_string(buffer->data());
                std::cout << self->getFormattedTime() << " Client -> Backend: Payload size in buffer: " << buffer->size() << " bytes. received from client websocket:" << str_received << "\n";

                // Set message type for backend before writing
                self->backend_ws_->text(self->client_ws_->got_text());
                self->backend_ws_->binary(self->client_ws_->got_binary());

                // Write to backend
                std::cout << self->getFormattedTime() << " Client -> Backend: Awaiting async_write to backend...\n";
                self->backend_ws_->async_write(buffer->data(),
                    [self, buffer](beast::error_code ec_write, std::size_t bytes_written) mutable {
                        if (ec_write) {
                            std::cerr << self->getFormattedTime() << " Error writing to backend: " << ec_write.message() << "\n";
                            self->close_connections();
                            return;
                        }
                        std::cout << self->getFormattedTime() << " Client -> Backend: Wrote " << bytes_written << " bytes.\n";
                        // Continue reading from client
                        self->read_from_client_websocket();
                    }
                );
            }
        );
    }

    // Read from backend and write to client
    void read_from_backend_websocket() {
        auto self = shared_from_this();
        if (!client_ws_ || !client_ws_->is_open() || !backend_ws_ || !backend_ws_->is_open()) {
            std::cout << getFormattedTime() << " Client or Backend WebSocket closed in read_from_backend. Exiting loop.\n";
            close_connections();
            return;
        }

        auto buffer = std::make_shared<beast::flat_buffer>();
        std::cout << getFormattedTime() << " Backend -> Client: Awaiting async_read from backend...\n";
        backend_ws_->async_read(*buffer,
            [self, buffer](beast::error_code ec_read, std::size_t bytes_transferred) mutable {
                if (ec_read == websocket::error::closed || ec_read == asio::error::eof || ec_read == asio::error::operation_aborted) {
                    std::cout << self->getFormattedTime() << " Backend WebSocket closed or aborted in read_from_backend. " << ec_read.message() << "\n";
                    self->close_connections();
                    return;
                }
                if (ec_read) {
                    std::cerr << self->getFormattedTime() << " Error reading from backend: " << ec_read.message() << "\n";
                    self->close_connections();
                    return;
                }

                // Add this line to see the actual size of the payload available in the buffer
                std::string str_received = beast::buffers_to_string(buffer->data());
                std::cout << self->getFormattedTime() << " Backend -> Client: Payload size in buffer: " << buffer->size() << " bytes. received from backend websocket:" << str_received << "\n";

                std::cout << self->getFormattedTime() << " Backend -> Client: Read " << bytes_transferred
                            << " bytes. Is Text: " << self->backend_ws_->got_text() << ", Is Binary: " << self->backend_ws_->got_binary() << "\n";

                // Set message type for client before writing
                self->client_ws_->text(self->backend_ws_->got_text());
                self->client_ws_->binary(self->backend_ws_->got_binary());

                // Write to client
                std::cout << self->getFormattedTime() << " Backend -> Client: Awaiting async_write to client...\n";
                self->client_ws_->async_write(buffer->data(),
                    [self, buffer](beast::error_code ec_write, std::size_t bytes_written) mutable {
                        if (ec_write) {
                            std::cerr << self->getFormattedTime() << " Error writing to client: " << ec_write.message() << "\n";
                            self->close_connections();
                            return;
                        }
                        std::cout << self->getFormattedTime() << " Backend -> Client: Wrote " << bytes_written << " bytes.\n";
                        // Continue reading from backend
                        self->read_from_backend_websocket();
                    }
                );
            }
        );
    }

    // Centralized connection closing logic
    void close_connections() {
        beast::error_code ec;

        if (client_ws_ && client_ws_->is_open()) {
            std::cout << getFormattedTime() << "Closing client WebSocket...\n";
            client_ws_->close(websocket::close_code::normal, ec);
            if (ec && ec != websocket::error::closed && ec != asio::error::eof && ec != asio::error::operation_aborted) {
                std::cerr << getFormattedTime() << " Error closing client WebSocket: " << ec.message() << "\n";
            }
            client_ws_.reset();
        } else if (stream_.lowest_layer().is_open()) { // If client_ws_ wasn't created yet, close the underlying SSL stream
            std::cout << getFormattedTime() << "Closing underlying client SSL stream...\n";
            stream_.lowest_layer().shutdown(tcp::socket::shutdown_both, ec);
            stream_.lowest_layer().close(ec);
        }

        if (backend_ws_ && backend_ws_->is_open()) {
            std::cout << getFormattedTime() << "Closing backend WebSocket...\n";
            backend_ws_->close(websocket::close_code::normal, ec);
            if (ec && ec != websocket::error::closed && ec != asio::error::eof && ec != asio::error::operation_aborted) {
                std::cerr << getFormattedTime() << " Error closing backend WebSocket: " << ec.message() << "\n";
            }
            backend_ws_.reset();
        } else if (backend_socket_.is_open()) { // If backend_ws_ wasn't created yet, close the underlying TCP socket
            std::cout << getFormattedTime() << "Closing underlying backend TCP socket...\n";
            backend_socket_.shutdown(tcp::socket::shutdown_both, ec);
            backend_socket_.close(ec);
        }
    }
};

// The server class to accept new connections
class server {
public:
    server(asio::io_context& ioc, unsigned short port, KedisConfig& config)
        :   ioc_(ioc),
            acceptor_(ioc, {tcp::v4(), port}),
            ssl_ctx_(ssl::context::tlsv12_server), // Use TLSv1.2 server context
            config_(config)
    {        
        std::string bind = config_.get("bind");        
        std::string CERT_FILE = config_.get("cert_file");
        std::string KEY_FILE = config_.get("key_file");

        // Configure SSL context (replace with your actual certificate and key paths)
        ssl_ctx_.set_options(
            ssl::context::default_workarounds |
            ssl::context::no_sslv2 |
            ssl::context::no_sslv3 |
            ssl::context::single_dh_use);

        // Replace with your actual certificate and key files
        // For testing, you can generate self-signed certs with mkcert:
        // mkcert localhost 127.0.0.1 ::1
        // Then use localhost+3.pem and localhost+3-key.pem (or similar names)
        // #ifdef _WIN32
        // const std::string CERT_FILE = "c:\\Users\\Admin\\v.host.pem";
        // const std::string KEY_FILE = "c:\\Users\\Admin\\v.host-key.pem";
        // #else
        // const std::string CERT_FILE = "/home/tom/v.host.pem";
        // const std::string KEY_FILE = "/home/tom/v.host-key.pem";
        // #endif
        try {
            ssl_ctx_.use_certificate_chain_file(CERT_FILE); // Path to your certificate chain
            ssl_ctx_.use_private_key_file(KEY_FILE, ssl::context::pem); // Path to your private key
        } catch (const boost::system::system_error& e) {
            std::cerr << "SSL context setup error: " << e.what() << "\n";
            exit(EXIT_FAILURE);
        }

        do_accept(); // Start accepting connections
    }

private:
    asio::io_context& ioc_;
    tcp::acceptor acceptor_;
    ssl::context ssl_ctx_;
    KedisConfig& config_;

    void do_accept() {
        acceptor_.async_accept(
            [this](beast::error_code ec, tcp::socket socket) {
                if (!ec) {
                    // Create a new session for each accepted client
                    std::make_shared<session>(ioc_, ssl_ctx_, std::move(socket), config_)->start();
                } else {
                    std::cerr << "Accept error: " << ec.message() << "\n";
                }
                do_accept(); // Continue accepting new connections
            }
        );
    }
};

int main() {
    try {
        asio::io_context ioc;
        KedisConfig config("reverse_proxy.conf");
        std::string str_port = config.get("port");
        uint32_t port = std::stoul(str_port);
        server s(ioc, port, config);
        std::string BACKEND_HOST_HTTP = config.get("backend_host_http");
        std::string BACKEND_PORT_HTTP = config.get("backend_port_http");
        std::string BACKEND_HOST_WS = config.get("backend_host_ws");
        std::string BACKEND_PORT_WS = config.get("backend_port_ws");
        std::cout << "Server listening on port " << port << " (SSL/TLS).\n";
        std::cout << "Proxying to HTTP backend: " << BACKEND_HOST_HTTP << ":" << BACKEND_PORT_HTTP << "\n";
        std::cout << "Proxying to WebSocket backend: " << BACKEND_HOST_WS << ":" << BACKEND_PORT_WS << "\n";

        ioc.run(); // Run the Boost.Asio event loop
    } catch (const std::exception& e) {
        std::cerr << "Exception: " << e.what() << "\n";
    }
    return 0;
}
